package prefix;

/**
 * 快速计算字符串
 *  出现次数
 *  指定前（后）缀出现次数
 */
public class TrieByClass {

    public static void main(String[] args) {
        TrieByClass trie = new TrieByClass();
        trie.insert("apple");
        trie.insert("apple");
        trie.insert("appcd");
        trie.delete("apple");
        System.out.println(trie.prefixNumber("ap"));
    }

    public class Node {
        // 经过当前节点的字符串个数
        public int passCount;
        // 以当前节点为终点的字符串个数
        public int endCount;

        public Node[] children;

        public Node() {
            passCount = 0;
            endCount = 0;
            children = new Node[26];
        }
    }

    private final Node root;

    public TrieByClass() {
        root = new Node();
    }

    public void insert(String word) {
        Node node = root;
        node.passCount++;

        for (int i = 0; i < word.length(); i++) {
            int path = word.charAt(i) - 'a';
            if (node.children[path] == null) {
                node.children[path] = new Node();
            }
            node = node.children[path];
            node.passCount++;
        }
        node.endCount++;
    }

    public int search(String word) {
        Node node = root;
        for (int i = 0; i < word.length(); i++) {
            int path = word.charAt(i) - 'a';
            if (node.children[path] == null) {
                return 0;
            }
            node = node.children[path];
        }
        return node.endCount;
    }

    public void delete(String word) {
        if (search(word) == 0) {
            return;
        }

        Node node = root;
        node.passCount--;

        for (int i = 0; i < word.length(); i++) {
            int path = word.charAt(i) - 'a';
            // 如果为 1 需要移除
            if (node.children[path].passCount == 1) {
                node.children[path] = null;
                return;
            }
            node = node.children[path];
            node.passCount--;
        }
        node.endCount--;
    }

    public int prefixNumber(String prefix) {
        Node node = root;
        for (int i = 0; i < prefix.length(); i++) {
            int path = prefix.charAt(i) - 'a';
            if (node.children[path] == null) {
                return 0;
            }
            node = node.children[path];
        }
        return node.passCount;
    }

}
